#!/usr/bin/env python
# coding: utf-8

# 搭建并训练CNN神经网络
import random
import os
import pylab
from PIL import Image
import numpy as np
from matplotlib.pyplot import imshow
import matplotlib.pyplot as plt
import sys
from matplotlib import pyplot
from keras.utils import to_categorical
from keras.models import Sequential
from keras.layers import Conv2D
from keras.layers import MaxPool2D
from keras.layers import Dense
from keras.layers import Flatten
from keras.layers import Dropout
from keras.optimizers import SGD
from keras.preprocessing.image import ImageDataGenerator
from keras.models import load_model
import tensorflow as tf
tf.compat.v1.logging.set_verbosity(tf.compat.v1.logging.ERROR)


def define_cnn_model():
    # 使用Sequential序列模型
    model = Sequential()
    # 卷积层
    model.add(Conv2D(32, (3, 3), activation="relu", padding="same", input_shape=(200, 200, 3)))  # 第一层即为卷积层，要设置输入进来图片的样式  3是颜色通道个数
    # 最大池化层
    model.add(MaxPool2D((2, 2)))  # 池化窗格
    # Flatten层
    model.add(Flatten())
    # 退出层 减少神经网络过拟合的结构
    model.add(Dropout(0.5))
    # 全连接层
    model.add(Dense(128, activation="relu"))  # 128为神经元的个数
    model.add(Dense(1, activation="sigmoid"))
    # 编译模型
    opt = SGD(learning_rate=0.0001, momentum=0.9)  # 随机梯度
    model.compile(optimizer=opt, loss="binary_crossentropy", metrics=["accuracy"])
    return model


def train_cnn_model():
    # 实例化模型
    model = define_cnn_model()
    # 创建图片生成器
    # ImageDataGenerator图像生成加数据增强
    datagen = ImageDataGenerator(rescale=1.0/255.0)
    train_it = datagen.flow_from_directory('E:\\Cat_And_Dog\\train1\\train\\',
        class_mode="binary",  # 该参数决定了返回的标签数组的形式，二分类问题"binary"返回1D的二值标签
        batch_size=64,
        target_size=(200, 200))  # batch_size:一次拿出多少张照片 targe_size:将图片缩放到一定比例
    # 训练模型
    model.fit(train_it,
              steps_per_epoch=len(train_it),
              epochs=150,  # 学习次数
              verbose=1)  # 当verbose=1时，带进度条的输出日志信息
    model.save("my_model.h5")


train_cnn_model()

# 识别train文件夹中的猫狗
model_path = "my_model.h5"
model = load_model(model_path)
plt.rcParams['font.sans-serif'] = ['SimHei']  # 设置字体


def read_random_image():
    folder = r'D:\\Cat_And_Dog\\train\\'
    file_path = folder + random.choice(os.listdir(folder))
    pil_im = Image.open(file_path, 'r')
    return pil_im


def get_predict(pil_im, model):
    # 首先更改图片的大小
    name = ''
    pil_im = pil_im.resize((200, 200))
    # 将格式转为numpy array格式
    array_im = np.asarray(pil_im)
    # array_im = array_im.resize((4,4))
    array_im = array_im[np.newaxis, :]
    # 对图像检测
    result = model.predict([[array_im]])
    if result[0][0] > 0.5:
        name = "它是狗！"
        print("结果为", '%d' % result[0][0])
        print("预测结果是：狗")
    else:
        name = "它是猫！"
        print("结果为", '%d' % result[0][0])
        print("预测结果是：猫")
    return name


pil_im = read_random_image()
imshow(np.asarray(pil_im))  # 显示图像
plt.title(get_predict(pil_im, model))
pylab.show()


# 识别NCEPUcats中的华电猫狗
def read_random_image():
    folder = r'D:\\Cat_And_Dog\\NCEPUcats\\'
    file_path = folder + random.choice(os.listdir(folder))
    pil_im = Image.open(file_path, 'r')
    return pil_im


def get_predict(pil_im, model):
    # 首先更改图片的大小
    name = ''
    pil_im = pil_im.resize((200, 200))
    # 将格式转为numpy array格式
    array_im = np.asarray(pil_im)
    # array_im = array_im.resize((4,4))
    array_im = array_im[np.newaxis, :]
    # 对图像检测
    result = model.predict([[array_im]])
    if result[0][0] > 0.5:
        name = "它是狗！"
        print("结果为", '%d' % result[0][0])
        print("预测结果是：狗")
    else:
        name = "它是猫！"
        print("结果为", '%d' % result[0][0])
        print("预测结果是：猫")
    return name


pil_im = read_random_image()
imshow(np.asarray(pil_im))
plt.title(get_predict(pil_im, model))
pylab.show()

